import time

import mayavi.mlab as mlab
import numpy as np

from util.point_transform import *
from util.seg_util import *
import util.visualize_utils as V
from matplotlib.cm import get_cmap

cmap = get_cmap('jet')
# 渲染时不需要显示图片
# mlab.options.offscreen = True

# # 输出图片的高宽
# output_img_h = 1440  # 1080
# output_img_w = 2560  # 1920
# size = (output_img_h, output_img_w)
# d = [0, 0, 38]
# theta = [0, -math.pi * (15.0 / 18.0), math.pi / 2]
# vis_point_radius = 1
# Bbox = None

# type = {
#     0: "Car",
#     1: "Bus",
#     2: 'construction_vehicle',
#     3: 'trailer',
#     4: 'Truck',
#     5: 'Bicycle',
#     6: 'Motorcycle',
#     7: 'Pedestrian',
#     8: 'Traffic_cone',
#     9: 'barrier'
# }
scene_param={
    "azimuth":176,
    "elevation":1e-5,
    "distance":350.0,
    "roll":0,
    "focalpoint":[150, 0, 0],
    "size":[2560,1440]
}


# def LabelPoints2ImagebyProjection(out_lidar_path_name, label_file, pt_file):
#     name,box=get_box_label(label_file)
#     name = np.array(name)
#     pt = get_lidar(pt_file)
#     label, corners_lidar, seg_types = Get_seglabel(name, box, pt)
#     seg_types = seg_types
#     pt = np.concatenate([pt[:, :3], seg_types.reshape(-1, 1), pt[:, 3].reshape(-1, 1)], axis=-1)
#     vis_img = VisualizePointsClass(pt, output_img_h, output_img_w, d=d, theta=theta,
#                                    vis_point_radius=vis_point_radius,
#                                    Bbox=corners_lidar,label=label)
#     return vis_img
    # import pdb
    # pdb.set_trace()
    # SaveVisImg(vis_img, './'+pts[:-4], out_lidar_path_name)

def LabelPoints2ImagebyMaya(out_lidar_path_name, pt_file, label_file=None, gtlabel=None, label_flag=1,
                            azimuth=176, elevation=70, distance=300.0, roll=90.0, focalpoint=[0, 0, 0],size=[2560,1440]):
    '''
     :out_lidar_path_name: str. The out lidar image path and name.

     :label_file: str. The label file path and name.

     :gtlabel: str. The real label file path and name.

     :label_flag: {1,2,3}.1:only label_file; 2:only gtlabel; 3:label_file and gtlabel

     :pt_file: str. The point file path and name.

     :azimuth: float, optional. The azimuthal angle (in degrees, 0-360),
        i.e. the angle subtended by the position vector on a sphere
        projected on to the x-y plane with the x-axis.

     :elevation: float, optional. The zenith angle (in degrees, 1e-10-180),
        i.e. the angle subtended by the position vector and the z-axis.

     :distance: float or 'auto', optional.
        A positive floating point number representing the distance from
        the focal point to place the camera. New in Mayavi 3.4.0: if
        'auto' is passed, the distance is computed to have a best fit of
        objects in the frame.

     :focalpoint: array_like or 'auto', optional.
        An array of 3 floating point numbers representing the focal point
        of the camera. New in Mayavi 3.4.0: if 'auto' is passed, the
        focal point is positioned at the center of all objects in the
        scene.

     :roll: float, optional
        Controls the roll, ie the rotation of the camera around its axis.

    :param size: The size of outputhe picture
    :return: None
    '''

    if label_flag == 1:
        # 只有检测目标
        name, box = get_box_label(label_file)
        name = np.array(name)
    elif label_flag == 2:
        # 只有真值目标
        name, box = get_box_label(gtlabel)
        name = np.array(name)
    else:
        # 既有真值目标也有检测目标
        name_pr, box_pr = get_box_label(label_file)
        for i in range(name_pr.__len__()):
            name_pr[i] = 'pr_label'
        name_pr = np.array(name_pr)
        # name_pr[:] = 'pedestrians'
        name_gt, box_gt = get_box_label(gtlabel)
        for i in range(name_gt.__len__()):
            name_gt[i] = 'gt_label'
        name_gt = np.array(name_gt)
        if (box_pr.__len__()>0) and box_gt.__len__()<=0:
            box = np.array(box_pr)
            name = np.array(name_pr)
        elif (box_pr.__len__()<=0) and box_gt.__len__()>0:
            box = np.array(box_gt)
            name = np.array(name_gt)
        elif (box_pr.__len__()>0) and box_gt.__len__()>0:
            box = np.concatenate([box_pr, box_gt])
            name = np.concatenate([name_pr, name_gt])
        else:
            return
    # pass
    pt = get_lidar(pt_file)
    # print(name)
    label, corners_lidar, seg_types = Get_seglabel(name, box, pt)
    pt = np.concatenate([pt[:, :3], seg_types.reshape(-1, 1), pt[:, 3].reshape(-1, 1)], axis=-1)

    # # 垂直映射
    # pt[:, 2] = 0
    # box[:, 2] = 0
    # box[:, 5] = 1e-10

    # # 计算需要绘图的类别数
    # color_classes = np.unique(name)
    # # 根据出现频率排序
    # color_classes_dict = []
    # for color_class in color_classes:
    #     color_classes_dict.append(sum(name == color_class))
    # index = np.argsort(color_classes_dict)
    # color_classes = color_classes[index]
    # print(color_classes)
    color_classes = np.array(['gt_label','pr_label'])
    # 作类别颜色映射
    color_maps = {}
    color_index = [120,180]
    for i in range(color_classes.__len__()):
        # color_index = int((i + 1) * cmap.N / color_classes.__len__())
        color_maps[color_classes[i]] = cmap(color_index[i])[0:3]
    color_maps['point'] = cmap(0)[0:3]

    # 给出映射表
    # # print(color_maps)
    # def plot_colormap(color_maps):
    #     fig, ax = plt.subplots(figsize=(6, 2))
    #     ax.imshow([color_maps[i] for i in color_maps], extent=[0, 10, 0, 1])
    #     ax.set_xticks([])
    #     ax.set_yticks([])
    #     ax.set_title(cmap_name)

    # vis_img = VisualizePointsClass(pt, output_img_h, output_img_w, d=d, theta=theta,
    #                                vis_point_radius=vis_point_radius,
    #                                Bbox=corners_lidar,label=label)
    # import pdb
    # pdb.set_trace()
    bgcolor = (0, 0, 0)
    fgcolor = (1, 1, 1)
    fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=fgcolor, engine=None, size=size)
    fig = V.draw_scenes(pt, gt_boxes=None, ref_boxes=box, ref_scores=None, ref_labels=name, size=size, fig=fig,
                        color_maps=color_maps)

    for color_class in color_classes:
        # 车辆位置从几何中点变换到车尾中点
        box_temp = box[name == color_class]
        l = box_temp[:, 3].reshape([1, 1, box_temp.__len__()])

        phi = box_temp[:, 6] * 180 / np.pi
        # phi = phi % 360
        # phi[phi > 270] = phi[phi > 270] - 360
        # phi[(phi > 90) & (phi <= 270)] = phi[(phi > 90) & (phi <= 270)] - 180
        #
        # # 车辆位置从几何中点变换到车位中点
        sin_phi = np.sin(phi / 180 * np.pi)
        cos_phi = np.cos(phi / 180 * np.pi)
        #
        # box_temps = np.array([cos_phi, sin_phi]).T * (-1 / 2) * np.array([box_temp[:, 3], box_temp[:, 3]]).T + box_temp[
        #                                                                                                        :, 0:2]

        x = box_temp[:, 0].reshape([1, 1, box_temp.__len__()])
        y = box_temp[:, 1].reshape([1, 1, box_temp.__len__()])
        z = box_temp[:, 2].reshape([1, 1, box_temp.__len__()])
        u = x / x * cos_phi.reshape([1, 1, box_temp.__len__()]) * l
        v = y / y * sin_phi.reshape([1, 1, box_temp.__len__()]) * l
        w = z
        w[:] = 0
        mlab.quiver3d(x, y, z, u, v, w, figure=fig, color=color_maps[color_class], scale_factor=1)
    # fig = visualize_pts(pt, show_intensity=True, size=size, fig=fig)
    mlab.view(azimuth=azimuth, elevation=elevation, distance=distance, roll=roll, focalpoint=focalpoint, figure=fig)
    # mlab.view(azimuth=176, elevation=70, distance=300.0, roll=90.0, focalpoint=[0, 0, 0],figure=fig)
    # time.sleep(1)
    # 看效果
    mlab.show(stop=True)
    # mlab.savefig(out_lidar_path_name, figure=fig)
    # mlab.close(all=True)